import math
import numpy as np
from gensim.models import Word2Vec
import jieba
from sklearn.cluster import KMeans
from collections import defaultdict
# 先训练词向量 根据词向量计算
def train_model(path):
    sentence = set() # 解析出来的每一行
    with open(path,'r',encoding='utf-8') as f:
        for line in f:
            sentence.add(" ".join(jieba.lcut(line.strip())))
    sentences_cut = [ a.split() for a in sentence]
    return Word2Vec(sentences_cut,vector_size=128,sg=1) # 训练模型

# 根据训练好的词向量模型将句子转为向量
def to_vector(model,sentences):
    vectors = []
    for sentence in sentences:
        vector = np.zeros(model.vector_size)
        a = jieba.lcut(sentence)
        for word in a:
            if word in model.wv:
                vector+=model.wv[word]
            else:
                vector+=np.zeros(model.vector_size)
        vectors.append(vector/len(a))
    return vectors

def main():
    model = train_model("素材/corpus.txt")
    sentences = []
    with open("素材/titles.txt", 'r', encoding='utf-8') as f:
        sentences = [ i.strip() for i in f]
    vectors = to_vector(model,sentences)
    # 开始聚类
    n_clusters = int(math.sqrt(len(sentences)))
    kmeans = KMeans(n_clusters=n_clusters)
    kmeans.fit(vectors)
    # 聚类结束
    # 现在我要求每类的平均距离 并且找出平均距离最小的三类中 离类中心最近的10个句子
    k=defaultdict(dict)
    for index,(vector,label) in enumerate(zip(vectors,kmeans.labels_)):
        k[label][index]=np.sqrt(np.sum(np.square(vector-kmeans.cluster_centers_[label])))
    # 2. 计算每类的平均距离，并按平均距离排序
    # 存储：{类别: 平均距离, ...}
    label_avg_distance = {}
    for label, index_dist in k.items():
        # 该类所有样本的距离列表
        distances = list(index_dist.values())
        # 计算平均距离
        avg_dist = np.mean(distances)
        label_avg_distance[label] = avg_dist

    # 按平均距离从小到大排序（得到平均距离最小的类别在前）
    sorted_labels = sorted(label_avg_distance.keys(), key=lambda x: label_avg_distance[x])
    top3_labels = sorted_labels[:3]

    top10_in_top3 = {}
    for clus in top3_labels:
        a = sorted(k[clus].items(), key=lambda x: x[1])
        top10 = a[:10]
        top10_in_top3[clus] = [label[0] for label in top10]

    # 打印一下结果
    for index,clus in enumerate(top3_labels):
        print("类内距离最近第%d名是第%d类"%(index,clus))
        for label in top10_in_top3[clus]:
            print(sentences[label])
        print()



if __name__ == '__main__':
    main()
